-- 基于CAS的令牌桶算法Lua脚本实现
-- 确保在分布式环境下的原子性操作，并支持等待策略
-- 
-- KEYS[1]: 令牌桶的key前缀
-- ARGV[1]: 最大令牌数
-- ARGV[2]: 令牌生成速率（每秒）
-- ARGV[3]: 当前时间戳（毫秒）
-- ARGV[4]: 请求令牌数量（默认为1）
-- ARGV[5]: 版本号（用于CAS操作，可选）

-- 获取参数
local key = KEYS[1]
local tokensKey = key .. ":tokens"
local timestampKey = key .. ":timestamp"
local versionKey = key .. ":version"

local maxTokens = tonumber(ARGV[1])
local tokensPerSecond = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local requestTokens = tonumber(ARGV[4]) or 1
local clientVersion = ARGV[5] and tonumber(ARGV[5]) or nil

-- 获取当前令牌数、上次更新时间和版本号
local tokensStr = redis.call('get', tokensKey)
local timestampStr = redis.call('get', timestampKey)
local versionStr = redis.call('get', versionKey)

-- 如果key不存在，初始化为满额令牌
local currentTokens = maxTokens
local lastRefillTime = now
local currentVersion = 1

if tokensStr then
    currentTokens = tonumber(tokensStr)
    lastRefillTime = tonumber(timestampStr)
    currentVersion = versionStr and tonumber(versionStr) or 1
end

-- 如果提供了版本号，进行CAS检查
if clientVersion and clientVersion ~= currentVersion then
    -- 版本不匹配，CAS失败
    -- 返回格式: [0, 当前令牌数, 下一个令牌预计可用时间, 当前版本号]
    local nextTokenTime = 0
    if currentTokens < requestTokens then
        -- 计算下一个令牌可用的时间
        local tokensNeeded = requestTokens - currentTokens
        nextTokenTime = math.ceil((tokensNeeded / tokensPerSecond) * 1000)
    end
    return {0, currentTokens, nextTokenTime, currentVersion}
end

-- 计算从上次更新到现在应该新增的令牌数
local timePassed = math.max(0, now - lastRefillTime)
local newTokens = timePassed / 1000.0 * tokensPerSecond

-- 更新当前令牌数，不超过最大令牌数
currentTokens = math.min(maxTokens, currentTokens + newTokens)

-- 增加版本号
currentVersion = currentVersion + 1

-- 如果令牌数足够，则获取成功
if currentTokens >= requestTokens then
    -- 消耗令牌
    currentTokens = currentTokens - requestTokens
    
    -- 更新令牌数、时间戳和版本号
    redis.call('set', tokensKey, tostring(currentTokens))
    redis.call('set', timestampKey, tostring(now))
    redis.call('set', versionKey, tostring(currentVersion))
    
    -- 设置过期时间，避免长期占用Redis内存
    local expireTime = math.ceil(maxTokens / tokensPerSecond) * 2
    redis.call('expire', tokensKey, expireTime)
    redis.call('expire', timestampKey, expireTime)
    redis.call('expire', versionKey, expireTime)
    
    -- 返回格式: [1, 当前令牌数, 0, 当前版本号]
    return {1, currentTokens, 0, currentVersion} -- 获取令牌成功
else
    -- 计算下一个令牌可用的时间
    local tokensNeeded = requestTokens - currentTokens
    local nextTokenTime = math.ceil((tokensNeeded / tokensPerSecond) * 1000)
    
    -- 更新时间戳和版本号，但不消耗令牌
    redis.call('set', tokensKey, tostring(currentTokens))
    redis.call('set', timestampKey, tostring(now))
    redis.call('set', versionKey, tostring(currentVersion))
    
    -- 设置过期时间
    local expireTime = math.ceil(maxTokens / tokensPerSecond) * 2
    redis.call('expire', tokensKey, expireTime)
    redis.call('expire', timestampKey, expireTime)
    redis.call('expire', versionKey, expireTime)
    
    -- 返回格式: [0, 当前令牌数, 下一个令牌预计可用时间, 当前版本号]
    return {0, currentTokens, nextTokenTime, currentVersion} -- 获取令牌失败，令牌不足
end